from collections.abc import Iterable, Iterator
from typing import Any, Callable, TypeVar

from zkl_aiutils_datasets.basics import Dataset, DatasetIterator, pause, resume
from zkl_aiutils_datasets.wrapping import wrap_dataset, wrap_dataset_iterator

AnySample = TypeVar('AnySample')


class FilteredIterator(DatasetIterator[AnySample]):
    def __init__(self, iterator: Iterator[AnySample] | Iterable[AnySample], filter_func: Callable[[AnySample], bool]):
        self._iterator = wrap_dataset_iterator(iterator)
        self._filter_func = filter_func

    def __next__(self) -> AnySample:
        while True:
            sample = next(self._iterator)
            if self._filter_func(sample):
                return sample

    def __pause__(self) -> int:
        return pause(self._iterator)


class FilteredDataset(Dataset[AnySample]):
    def __init__(self, dataset: Iterable[AnySample], filter_func: Callable[[AnySample], bool]):
        self._dataset = wrap_dataset(dataset)
        self._filter_func = filter_func

    @property
    def dataset(self) -> Dataset[AnySample]:
        return self._dataset

    @property
    def filter_func(self) -> Callable[[AnySample], bool]:
        return self._filter_func

    def __iter__(self) -> FilteredIterator[AnySample]:
        return FilteredIterator(iter(self._dataset), self._filter_func)

    def __resume__(self, state: Any | None = None, /) -> FilteredIterator[AnySample]:
        return FilteredIterator(resume(self._dataset, state), self._filter_func)
